package com.cxyjy.cfyuiapi.websocket;

import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.cxyjy.cfyuiapi.service.impl.QueryHistoryServiceImpl;
import com.cxyjy.common.util.JObjectUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArraySet;

import static com.cxyjy.common.constants.CfyuiConstants.WEBSOCKET_URL_STR;

/**
 * @Author: jony
 * @Date : 2024/11/18 22:06
 * @VERSION v1.0
 */
@Slf4j
@Component
@ServerEndpoint("/ws")
public class WebSocket {

    //用来记录当前在线连接数。增加减少连接数是线程同步的
    private static int onlineCount = 0;

    //每个连接通过他来传递消息：javax.websocket.Session
    private Session session;

    //key:clientId value:Session
    private static final Map<String, Session> clientSessions = new HashMap<>();
    //每创建一个连接就会生成一个WebScoket,用CopyOnWriteArraySet进行保存
    private static CopyOnWriteArraySet<WebSocket> webSockets = new CopyOnWriteArraySet<>();

    //建立连接时自动调用
    @OnOpen
    public void onOpen(Session session) throws URISyntaxException, InterruptedException {
        this.session = session;
        session.getUserProperties().put("websocketInstance",this);
        Map<String, List<String>> requestParameterMap = session.getRequestParameterMap();
        //前端new的websocket url是"ws://baseurl:port/ws?clientId=********"
        if (requestParameterMap.containsKey("clientId")) {
            String clientId = requestParameterMap.get("clientId").get(0);
            //将客户端id和Session对象储存在clientSessions当中
            clientSessions.put(clientId, session);
            String uriStr = WEBSOCKET_URL_STR + "?clientId=" + clientId;
            URI workFlowUri = new URI(uriStr);
            //前端发送websocket连接请求，此时后端作为客户端在这里建立websocket连接，连接到comfyUI
            CfuiWebsocketClient client = new CfuiWebsocketClient(workFlowUri);
            client.connectBlocking(); // 阻塞直到连接建立或失败
        }
        addOnlineCount();
        webSockets.add(this);
        log.info("建立连接");
        sendMessage("建立连接！");
    }

    //连接关闭时自动调用
    @OnClose
    public void onClose() {
        webSockets.remove(this);
        subOnlineCount();
        log.info("断开连接");
    }

    //收到消息时自动调用
    @OnMessage
    public void onMessage(String msg, Session session) {
        try {
            Map<String, Object> msgMap = JObjectUtil.toMap(JSONObject.parse(msg));
            String type = JObjectUtil.stringValue(msgMap.get("type"));
            log.info(msg);
//            if (type.equals("progress")) {
//
//                int stepsValue = JObjectUtil.objectToInt(msgMap.get("value"));
//                int stepsMax = JObjectUtil.objectToInt(msgMap.get("max"));
//                Map<String,Object> replyMsgMap = new HashMap<>();
//                msgMap.put("msg","换脸工作流正在进行中，已经完成"+stepsValue/stepsMax*100+"%");
//                sendMessageToClient(clientId,"");
//                if (stepsValue == stepsMax) {
//                    closeSession(session);
//                }
//            }

        } catch (JSONException | ClassCastException e) {
            sendMessage("收到来自comfyui服务器的信息格式不正确！");
        }
        log.info("收到客户端发来的消息:" + msg);
    }

    //服务器发送消息的方法，给特定的发
    public void sendMessageToClient(String clientId, String message) {
        Session session = clientSessions.get(clientId);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                System.err.println("Failed to send message to client ID " + clientId + ": " + e.getMessage());
            }
        } else {
            System.out.println("No open session for client ID " + clientId);
        }
    }

    //给所有连接的websocket发送
    public void sendMessage(String msg) {

//        for (WebSocket webSocket : webSockets) {
//            log.info("{}发送消息：{}", webSocket, msg);
//            try {
//                //发送消息
//                webSocket.session.getBasicRemote().sendText(msg);
//            } catch (IOException e) {
//                e.printStackTrace();
//            }
//        }
    }

    public static Session getWebsocketSession(String promptId) throws IOException {
        Map<String,Map<String, List<Object>>> queryhistory =(Map) (new QueryHistoryServiceImpl().queryhistory(promptId));
        List<Object> promptObj = queryhistory.get(promptId).get("prompt");
        Object o = promptObj.get(3);
        Map<String, Object> promptMap = JObjectUtil.toMap(o);
        String clientId = JObjectUtil.stringValue(promptMap.get("client_id"));
        return clientSessions.get(clientId);
    }

    public static void closeSession(Session session) {
        try {
            if (session != null && session.isOpen()) {
                session.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "正常关闭"));
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        WebSocket.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        WebSocket.onlineCount--;
    }

}


//    在别的地方调用sendMessage() 发送消息
//    {
//        webSocket.sendMessage("新消息");
//    }
